{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Non-correspondences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/yoyee/Documents/pytorch-dense-correspondence/dense_correspondence\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style>.container { width:100% !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "module_path = os.path.abspath(os.path.join('..'))\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)\n",
    "print(module_path)\n",
    "%matplotlib inline\n",
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "# import dense_correspondence_manipulation.utils.utils as utils\n",
    "# utils.add_dense_correspondence_to_python_path()\n",
    "import correspondence_plotter\n",
    "import correspondence_finder\n",
    "# from dense_correspondence.dataset.spartan_dataset_masked import SpartanDataset\n",
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "tuple indices must be integers or slices, not tuple",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-99-2c0385012671>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m                 torch.tensor(np.array([1]), dtype=torch.float32))\n\u001b[1;32m      8\u001b[0m \u001b[0mmatches_a\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muv_a_matches\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mH\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0muv_a_matches\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m \u001b[0mmatches_a\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muvto1d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0muv_a_matches\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"uv_a_matches: \"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0muv_a_matches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"matches_a: \"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmatches_a\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-99-2c0385012671>\u001b[0m in \u001b[0;36muvto1d\u001b[0;34m(points, H)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0muvto1d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoints\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0;31m# assert points.dim == 2\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mpoints\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mH\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mpoints\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mW\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m30\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m40\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: tuple indices must be integers or slices, not tuple"
     ]
    }
   ],
   "source": [
    "def uvto1d(points, H):\n",
    "    # assert points.dim == 2\n",
    "    return points[:,0]*H + points[:,1]\n",
    "\n",
    "H, W = 30, 40\n",
    "uv_a_matches = (torch.tensor(np.array([1]), dtype=torch.float32), \n",
    "                torch.tensor(np.array([1]), dtype=torch.float32))\n",
    "matches_a = uv_a_matches[0]*H + uv_a_matches[1]\n",
    "matches_a = uvto1d(uv_a_matches, H)\n",
    "print(\"uv_a_matches: \", uv_a_matches)\n",
    "print(\"matches_a: \", matches_a)\n",
    "\n",
    "uv_b_matches = (torch.tensor(np.array([1]), dtype=torch.float32), \n",
    "                torch.tensor(np.array([1]), dtype=torch.float32))\n",
    "matches_b = uvto1d(uv_b_matches, H)\n",
    "\n",
    "img_b_shape = (H, W)\n",
    "img_a_shape = img_b_shape\n",
    "\n",
    "# image_a_pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.002017498016357422 seconds for non-matches\n",
      "torch.Size([1, 10])\n",
      "tensor(1.)\n",
      "tensor(4.)\n",
      "tensor(36.)\n",
      "tensor(28.)\n"
     ]
    }
   ],
   "source": [
    "# num_attempts = 5\n",
    "\n",
    "# img_a_index = dataset.get_random_image_index(scene)\n",
    "# img_a_rgb, img_a_depth, _, img_a_pose = dataset.get_rgbd_mask_pose(scene, img_a_index)\n",
    "\n",
    "# img_b_index = dataset.get_img_idx_with_different_pose(scene, img_a_pose, num_attempts=50)\n",
    "# img_b_rgb, img_b_depth, _, img_b_pose = dataset.get_rgbd_mask_pose(scene, img_b_index)\n",
    "\n",
    "# img_a_depth_numpy = np.asarray(img_a_depth)\n",
    "# img_b_depth_numpy = np.asarray(img_b_depth)\n",
    "\n",
    "# start = time.time()\n",
    "# uv_a, uv_b = correspondence_finder.batch_find_pixel_correspondences(img_a_depth_numpy, img_a_pose, \n",
    "#                                                                     img_b_depth_numpy, img_b_pose,\n",
    "#                                                                     num_attempts=num_attempts,\n",
    "#                                                                     device='CPU')\n",
    "\n",
    "\n",
    "start = time.time()\n",
    "# uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b, img_a_depth_numpy.shape, num_non_matches_per_match=10)\n",
    "uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b_matches, img_b_shape, num_non_matches_per_match=10, img_b_mask=None)\n",
    "print  (time.time() - start, \"seconds for non-matches\")\n",
    "if uv_b_non_matches is not None:\n",
    "    print (uv_b_non_matches[0].shape)\n",
    "\n",
    "    import torch\n",
    "    # This just checks to make sure nothing is out of bounds\n",
    "    print (torch.min(uv_b_non_matches[0]))\n",
    "    print (torch.min(uv_b_non_matches[1]))\n",
    "    print (torch.max(uv_b_non_matches[0]))\n",
    "    print (torch.max(uv_b_non_matches[1]))\n",
    "    \n",
    "#     fig, axes = correspondence_plotter.plot_correspondences_direct(img_a_rgb, img_a_depth_numpy, img_b_rgb, img_b_depth_numpy, uv_a, uv_b, show=False)\n",
    "#     uv_a_long = (torch.t(uv_a[0].repeat(3, 1)).contiguous().view(-1,1), torch.t(uv_a[1].repeat(3, 1)).contiguous().view(-1,1))\n",
    "#     uv_b_non_matches_long = (uv_b_non_matches[0].view(-1,1), uv_b_non_matches[1].view(-1,1) )\n",
    "#     correspondence_plotter.plot_correspondences_direct(img_a_rgb, img_a_depth_numpy, img_b_rgb, img_b_depth_numpy, uv_a_long, uv_b_non_matches_long, use_previous_plot=(fig,axes),\n",
    "#                                                   circ_color='r')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "non_matches_a:  tensor([[ 862.,  286., 1089.,  478.,  977.,  337.,  390.,  951.,   12.,  132.]])\n"
     ]
    }
   ],
   "source": [
    "uv_a_non_matches = correspondence_finder.create_non_correspondences(uv_a_matches, img_a_shape, num_non_matches_per_match=10, img_b_mask=None)\n",
    "uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b_matches, img_b_shape, num_non_matches_per_match=10, img_b_mask=None)\n",
    "\n",
    "non_matches_a = uvto1d(uv_a_non_matches, H)\n",
    "non_matches_b = uvto1d(uv_b_non_matches, H)\n",
    "\n",
    "print(\"non_matches_a: \", non_matches_a)\n",
    "\n",
    "# uv_b_non_matches # (u, v)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tensor(np.array([1]), dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loss_functions.pixelwise_contrastive_loss import PixelwiseContrastiveLoss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pltImshow(img):\n",
    "    from matplotlib import pyplot as plt\n",
    "    plt.imshow(img)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "D = 3\n",
    "# image_a = torch.tensor(np.random.rand(H, W), dtype=torch.float32)\n",
    "# image_b = torch.tensor(np.random.rand(H, W), dtype=torch.float32)\n",
    "\n",
    "image_a_pred = torch.tensor(np.random.rand(1, H*W, D), dtype=torch.float32)\n",
    "image_b_pred = torch.tensor(np.random.rand(1, H*W, D), dtype=torch.float32)\n",
    "\n",
    "# pltImshow(image_a.numpy())\n",
    "# pltImshow(image_b.numpy())\n",
    "# print(\"image_a_pred: \", ))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "match_loss:  tensor(0.1347)\n",
      "matches_a_descriptors:  torch.Size([1, 1, 1, 3])\n",
      "matches_b_descriptors:  torch.Size([1, 1, 1, 3])\n",
      "non_match_loss:  tensor([0.0000, 0.0000, 0.0000, 0.1419, 0.0010, 0.0422, 0.0000, 0.0000, 0.0000,\n",
      "        0.0832])\n",
      "num_hard_negatives:  4\n",
      "non_matches_a_descriptors:  torch.Size([10, 3])\n"
     ]
    }
   ],
   "source": [
    "alpha = 0.5\n",
    "\n",
    "match_loss, matches_a_descriptors, matches_b_descriptors = \\\n",
    "    PixelwiseContrastiveLoss.match_loss(image_a_pred, image_b_pred, matches_a.long(), matches_b.long())\n",
    "non_match_loss, num_hard_negatives, non_matches_a_descriptors, non_matches_b_descriptors = \\\n",
    "    PixelwiseContrastiveLoss.non_match_descriptor_loss(image_a_pred, image_b_pred, \n",
    "                                                   non_matches_a.long().squeeze(), non_matches_b.long().squeeze())\n",
    "\n",
    "\n",
    "print(\"match_loss: \", match_loss)\n",
    "print(\"matches_a_descriptors: \", matches_a_descriptors.shape)\n",
    "print(\"matches_b_descriptors: \", matches_b_descriptors.shape)\n",
    "\n",
    "print(\"non_match_loss: \", non_match_loss)\n",
    "print(\"num_hard_negatives: \", num_hard_negatives)\n",
    "print(\"non_matches_a_descriptors: \", non_matches_a_descriptors.shape)\n",
    "\n",
    "# loss, match_loss, non_match_loss = \\\n",
    "#     PixelwiseContrastiveLoss.get_triplet_loss(image_a_pred,\n",
    "#                                         image_b_pred,\n",
    "#                                         matches_a.long(),\n",
    "#                                         matches_b.long(),\n",
    "#                                         non_matches_a.long(),\n",
    "#                                         non_matches_b.long(), \n",
    "#                                         alpha = alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'self' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-86-ccee89f5ba03>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0muv_a_masked_long\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0muv_b_masked_non_matches_long\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcreate_non_matches\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmatches_1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmatches_2_masked_non_matches\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_masked_non_matches_per_match\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'self' is not defined"
     ]
    }
   ],
   "source": [
    "uv_a_masked_long, uv_b_masked_non_matches_long = \\\n",
    "        self.create_non_matches(matches_1, matches_2_masked_non_matches, self.num_masked_non_matches_per_match)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "homographies:  [[[1. 0. 0.]\n",
      "  [0. 1. 0.]\n",
      "  [0. 0. 1.]]]\n",
      "homographies_H:  [[[[1. 0. 0.]\n",
      "   [0. 1. 0.]\n",
      "   [0. 0. 1.]]]]\n"
     ]
    }
   ],
   "source": [
    "def scale_homography(H, shape, shift=(-1,-1)):\n",
    "    height, width = shape[0], shape[1]\n",
    "    trans = np.array([[2./width, 0., shift[0]], [0., 2./height, shift[1]], [0., 0., 1.]])\n",
    "    H_tf = np.linalg.inv(trans) @ H @ trans\n",
    "    return H_tf\n",
    "\n",
    "homographies = np.identity(3)[np.newaxis,:,:]\n",
    "height, width = 240, 320\n",
    "image_shape = np.array([height, width])\n",
    "\n",
    "homographies_H = np.stack([scale_homography(homographies, image_shape, shift=(-1,-1)) for H in homographies])\n",
    "\n",
    "print(\"homographies: \", homographies)\n",
    "print(\"homographies_H: \", homographies_H)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "coor_cells:  tensor([[ 0.,  0.],\n",
      "        [ 0.,  1.],\n",
      "        [ 0.,  2.],\n",
      "        ...,\n",
      "        [29., 37.],\n",
      "        [29., 38.],\n",
      "        [29., 39.]])\n",
      "coor_cells:  torch.Size([1200, 2])\n",
      "matches_a:  torch.Size([1200])\n",
      "matches_b:  torch.Size([1200])\n"
     ]
    }
   ],
   "source": [
    "H, W = 240, 320\n",
    "cell_size = 8\n",
    "Hc, Wc = H//cell_size, W//cell_size\n",
    "\n",
    "device = 'cpu'\n",
    "\n",
    "def get_coor_cells(Hc, Wc, cell_size, device='cpu'):\n",
    "    coor_cells = torch.stack(torch.meshgrid(torch.arange(Hc), torch.arange(Wc)), dim=2)\n",
    "    coor_cells = coor_cells.type(torch.FloatTensor).to(device)\n",
    "#     coor_cells = coor_cells * cell_size + cell_size // 2\n",
    "    ## coord_cells is now a grid containing the coordinates of the Hc x Wc\n",
    "    ## center pixels of the 8x8 cells of the image\n",
    "    return coor_cells.view(-1, 2)\n",
    "\n",
    "coor_cells = get_coor_cells(Hc, Wc, cell_size=cell_size, device=device)\n",
    "print(\"coor_cells: \", coor_cells)\n",
    "print(\"coor_cells: \", coor_cells.shape)\n",
    "\n",
    "def warp_coor_cells_with_homographies(coor_cells, homographies):\n",
    "    # warped_coor_cells = warp_points(coor_cells.view([-1, 2]), homographies, device)\n",
    "#     warped_coor_cells = normPts(coor_cells.view([-1, 2]), shape)\n",
    "    warped_coor_cells = torch.stack((warped_coor_cells[:,1], warped_coor_cells[:,0]), dim=1) # (y, x) to (x, y)\n",
    "    warped_coor_cells = warp_points(warped_coor_cells, homographies, device)\n",
    "\n",
    "    warped_coor_cells = torch.stack((warped_coor_cells[:, :, 1], warped_coor_cells[:, :, 0]), dim=2)  # (batch, x, y) to (batch, y, x)\n",
    "\n",
    "    shape_cell = torch.tensor([H//cell_size, W//cell_size]).type(torch.FloatTensor).to(device)\n",
    "    # warped_coor_mask = denormPts(warped_coor_cells, shape_cell)\n",
    "\n",
    "#     warped_coor_cells = denormPts(warped_coor_cells, shape)\n",
    "    return warped_coor_cells\n",
    "\n",
    "matches_a = uvto1d(coor_cells, Hc)\n",
    "print(\"matches_a: \", matches_a.shape)\n",
    "\n",
    "matches_b = matches_a\n",
    "print(\"matches_b: \", matches_b.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "match_loss:  tensor(0.4838)\n",
      "matches_a_descriptors:  torch.Size([1, 1200, 3])\n",
      "matches_b_descriptors:  torch.Size([1, 1200, 3])\n"
     ]
    }
   ],
   "source": [
    "match_loss, matches_a_descriptors, matches_b_descriptors = \\\n",
    "    PixelwiseContrastiveLoss.match_loss(image_a_pred, image_b_pred, matches_a.long(), matches_b.long())\n",
    "print(\"match_loss: \", match_loss)\n",
    "print(\"matches_a_descriptors: \", matches_a_descriptors.shape)\n",
    "print(\"matches_b_descriptors: \", matches_b_descriptors.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_non_matches(self, uv_a, uv_b_non_matches, multiplier):\n",
    "    \"\"\"\n",
    "    Simple wrapper for repeated code\n",
    "    :param uv_a:\n",
    "    :type uv_a:\n",
    "    :param uv_b_non_matches:\n",
    "    :type uv_b_non_matches:\n",
    "    :param multiplier:\n",
    "    :type multiplier:\n",
    "    :return:\n",
    "    :rtype:\n",
    "    \"\"\"\n",
    "    uv_a_long = (torch.t(uv_a[0].repeat(multiplier, 1)).contiguous().view(-1, 1),\n",
    "                 torch.t(uv_a[1].repeat(multiplier, 1)).contiguous().view(-1, 1))\n",
    "\n",
    "    uv_b_non_matches_long = (uv_b_non_matches[0].view(-1, 1), uv_b_non_matches[1].view(-1, 1))\n",
    "\n",
    "    return uv_a_long, uv_b_non_matches_long"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "descriptors:  torch.Size([1, 3, 30, 40])\n",
      "image_a_pred:  torch.Size([1, 1200, 3])\n",
      "image_a_pred:  torch.Size([1, 1200, 3])\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'normPts' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-117-e82ac0a8b694>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     60\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"descriptors: \"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdescriptors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     61\u001b[0m \u001b[0mdescriptors_warped\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mHc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mWc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 62\u001b[0;31m \u001b[0mdescriptor_loss_sparse\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdescriptors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdescriptors_warped\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhomographies\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-117-e82ac0a8b694>\u001b[0m in \u001b[0;36mdescriptor_loss_sparse\u001b[0;34m(descriptors, descriptors_warped, homographies, mask_valid, cell_size, device, descriptor_dist, lamda_d, num_matching_attempts, num_masked_non_matches_per_match, **config)\u001b[0m\n\u001b[1;32m     27\u001b[0m     \u001b[0mhomographies_H\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mscale_homography\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhomographies\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimage_shape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshift\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mH\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mhomographies\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 29\u001b[0;31m     \u001b[0muv_b_matches\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mwarp_coor_cells_with_homographies\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0muv_a\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhomographies\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     30\u001b[0m     \u001b[0mmatches_a\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muvto1d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0muv_a\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"matches_a: \"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmatches_a\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-105-4f84d5bd154e>\u001b[0m in \u001b[0;36mwarp_coor_cells_with_homographies\u001b[0;34m(coor_cells, homographies)\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mwarp_coor_cells_with_homographies\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcoor_cells\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhomographies\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0;31m# warped_coor_cells = warp_points(coor_cells.view([-1, 2]), homographies, device)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 21\u001b[0;31m     \u001b[0mwarped_coor_cells\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnormPts\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcoor_cells\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     22\u001b[0m     \u001b[0mwarped_coor_cells\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwarped_coor_cells\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwarped_coor_cells\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# (y, x) to (x, y)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0mwarped_coor_cells\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mwarp_points\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwarped_coor_cells\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhomographies\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'normPts' is not defined"
     ]
    }
   ],
   "source": [
    "def descriptor_loss_sparse(descriptors, descriptors_warped, homographies, mask_valid=None,\n",
    "                           cell_size=8, device='cpu', descriptor_dist=4, lamda_d=250,\n",
    "                           num_matching_attempts=1000, num_masked_non_matches_per_match=10, **config):\n",
    "    \"\"\"\n",
    "    consider batches of descriptors\n",
    "    :param descriptors:\n",
    "        Output from descriptor head\n",
    "        tensor [batch_size, descriptors, Hc, Wc]\n",
    "    :param descriptors_warped:\n",
    "        Output from descriptor head of warped image\n",
    "        tensor [batch_size, descriptors, Hc, Wc]\n",
    "    \"\"\"\n",
    "    batch_size, Hc, Wc = descriptors.shape[0], descriptors.shape[2], descriptors.shape[3]\n",
    "    \n",
    "    image_a_pred = descriptors.view(batch_size, -1, Hc*Wc).transpose(1,2)  # torch [batch_size, H*W, D]\n",
    "    print(\"image_a_pred: \", image_a_pred.shape)\n",
    "    image_b_pred = descriptors_warped.view(batch_size, -1, Hc*Wc).transpose(1,2)  # torch [batch_size, H*W, D]\n",
    "    print(\"image_a_pred: \", image_a_pred.shape)\n",
    "\n",
    "    # matches\n",
    "    ## sample points (implement sampling later)\n",
    "    ## warp points\n",
    "        # coor_cell\n",
    "    uv_a = get_coor_cells(Hc, Wc, cell_size)\n",
    "        # warp coor_cell\n",
    "#     homographies_H = \n",
    "    homographies_H = np.stack([scale_homography(homographies, image_shape, shift=(-1,-1)) for H in homographies])\n",
    "    \n",
    "    uv_b_matches = warp_coor_cells_with_homographies(uv_a, homographies)\n",
    "    matches_a = uvto1d(uv_a, H)\n",
    "    print(\"matches_a: \", matches_a.shape)\n",
    "    matches_b = uvto1d(uv_b_matches, H)\n",
    "    print(\"matches_b: \", matches_b.shape)\n",
    "    \n",
    "    # filtering!!!\n",
    "    \n",
    "    ## select from descriptor\n",
    "    ## calculate matches loss\n",
    "    match_loss, matches_a_descriptors, matches_b_descriptors = \\\n",
    "        PixelwiseContrastiveLoss.match_loss(image_a_pred, image_b_pred, matches_a.long(), matches_b.long())\n",
    "    \n",
    "    # non matches\n",
    "    ## sample non matches\n",
    "    img_b_shape = (Hc, Wc)\n",
    "    uv_b_non_matches = correspondence_finder.create_non_correspondences(uv_b_matches, img_b_shape, num_non_matches_per_match=10, img_b_mask=None)\n",
    "    \n",
    "    ## create_non_correspondences\n",
    "    uv_a_masked_long, uv_b_masked_non_matches_long = \\\n",
    "        self.create_non_matches(matches_1, matches_2_masked_non_matches, self.num_masked_non_matches_per_match)\n",
    "    \n",
    "    ## non matches loss\n",
    "    non_match_loss, num_hard_negatives, non_matches_a_descriptors, non_matches_b_descriptors = \\\n",
    "    PixelwiseContrastiveLoss.non_match_descriptor_loss(image_a_pred, image_b_pred, \n",
    "                                                   non_matches_a.long().squeeze(), non_matches_b.long().squeeze())\n",
    "    \n",
    "    loss = lamda_d*match_loss + non_match_loss\n",
    "    pass\n",
    "    \n",
    "descriptors = torch.tensor(np.random.rand(1, D, Hc, Wc), dtype=torch.float32)\n",
    "print(\"descriptors: \", descriptors.shape)\n",
    "descriptors_warped = torch.tensor(np.random.rand(1, D, Hc, Wc), dtype=torch.float32)\n",
    "descriptor_loss_sparse(descriptors, descriptors_warped, homographies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'torch.LongTensor'"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matches_a.long().type()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss_on_dataset(dcn, data_loader, loss_config, num_iterations=500,):\n",
    "    \"\"\"\n",
    "\n",
    "    Computes the loss for the given number of iterations\n",
    "\n",
    "    :param dcn:\n",
    "    :type dcn:\n",
    "    :param data_loader:\n",
    "    :type data_loader:\n",
    "    :param num_iterations:\n",
    "    :type num_iterations:\n",
    "    :return:\n",
    "    :rtype:\n",
    "    \"\"\"\n",
    "    dcn.eval()\n",
    "\n",
    "    # loss_vec = np.zeros(num_iterations)\n",
    "    loss_vec = []\n",
    "    match_loss_vec = []\n",
    "    non_match_loss_vec = []\n",
    "    counter = 0\n",
    "    pixelwise_contrastive_loss = PixelwiseContrastiveLoss(dcn.image_shape, config=loss_config)\n",
    "\n",
    "    batch_size = 1\n",
    "\n",
    "    for i, data in enumerate(data_loader, 0):\n",
    "\n",
    "        # get the inputs\n",
    "        data_type, img_a, img_b, matches_a, matches_b, non_matches_a, non_matches_b, metadata = data\n",
    "        data_type = data_type[0]\n",
    "\n",
    "        if len(matches_a[0]) == 0:\n",
    "            print (\"didn't have any matches, continuing\")\n",
    "            continue\n",
    "\n",
    "        img_a = Variable(img_a.cuda(), requires_grad=False)\n",
    "        img_b = Variable(img_b.cuda(), requires_grad=False)\n",
    "\n",
    "        if data_type == \"matches\":\n",
    "            matches_a = Variable(matches_a.cuda().squeeze(0), requires_grad=False)\n",
    "            matches_b = Variable(matches_b.cuda().squeeze(0), requires_grad=False)\n",
    "            non_matches_a = Variable(non_matches_a.cuda().squeeze(0), requires_grad=False)\n",
    "            non_matches_b = Variable(non_matches_b.cuda().squeeze(0), requires_grad=False)\n",
    "\n",
    "        # run both images through the network\n",
    "        image_a_pred = dcn.forward(img_a)\n",
    "        image_a_pred = dcn.process_network_output(image_a_pred, batch_size)\n",
    "\n",
    "        image_b_pred = dcn.forward(img_b)\n",
    "        image_b_pred = dcn.process_network_output(image_b_pred, batch_size)\n",
    "\n",
    "        # get loss\n",
    "        if data_type == \"matches\":\n",
    "            loss, match_loss, non_match_loss = \\\n",
    "                pixelwise_contrastive_loss.get_loss(image_a_pred,\n",
    "                                                    image_b_pred,\n",
    "                                                    matches_a,\n",
    "                                                    matches_b,\n",
    "                                                    non_matches_a,\n",
    "                                                    non_matches_b)\n",
    "\n",
    "\n",
    "\n",
    "            loss_vec.append(loss.data[0])\n",
    "            non_match_loss_vec.append(non_match_loss.data[0])\n",
    "            match_loss_vec.append(match_loss.data[0])\n",
    "\n",
    "\n",
    "        if i > num_iterations:\n",
    "            break\n",
    "\n",
    "    loss_vec = np.array(loss_vec)\n",
    "    match_loss_vec = np.array(match_loss_vec)\n",
    "    non_match_loss_vec = np.array(non_match_loss_vec)\n",
    "\n",
    "    loss = np.average(loss_vec)\n",
    "    match_loss = np.average(match_loss_vec)\n",
    "    non_match_loss = np.average(non_match_loss_vec)\n",
    "\n",
    "    return loss, match_loss, non_match_loss\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
